import streamlit as st
from dotenv import load_dotenv
load_dotenv() # .env添加智普api key
from zai import ZhipuAiClient
import base64

def llm_chat(message_list):
    # 初始化客户端
    client = ZhipuAiClient()

    # 转换消息格式
    msgs=[]
    for message in  message_list:
        content = message[1] # 聊天信息
        role = message[0] # 角色（ai、human）
        img_base = None
        if len(message) > 2:
            img_base = message[2]

        if role == 'ai':
            msgs.append({'role':'assistant','content':content})
        elif role == 'human':
            # 动态content内容
            content_list = [
                {
                "type":"text",
                "text":content
                }
            ]
            if img_base is not None:
                content_list.append({
                    "type":"image_url",
                    "image_url":{
                        "url":img_base
                    }
                })
            msgs.append({'role':'user','content':content_list})


    response = client.chat.completions.create(
        model = "glm-4V-flash",
        messages=msgs,
        temperature=0.6
    )
    return response.choices[0].message.content

# 初始化
if st.session_state.get('first_run',None) is None:

    # 模拟数据，聊天历史
    messages = [('human','你好!'),
            ('ai','你好！我是智能聊天助理，可以完成各种任务。'),
            ('human','你叫什么？'),
    ]
    st.session_state['messages'] = messages

    st.session_state['first_run'] = True

st.markdown('## 智能聊天助理')


# 结果值 = 组件（）
prompt = st.chat_input(
    "请输入问题或上传文件",
    accept_file=True, # 可以上传文件
    file_type=["jpg","jpeg","png"], # 支持上传文件类型
)


if prompt is not None:
    # 获取提交prompt和文件内容并转换base64格式
    prompt_text = prompt["text"] # 提交文本内容

    if len(prompt["files"]) > 0:
        img_base = base64.b64encode(prompt["files"][0].getvalue()).decode("utf-8") # 文件内容base64编码转换

        # 添加
        st.session_state['messages'].append(('human',prompt_text,img_base))

    else:
        st.session_state['messages'].append(('human',prompt_text))

    # 调用模型聊天函数
    resp = llm_chat(st.session_state['messages'])
    # 添加到messages
    st.session_state['messages'].append(('ai',resp))



for msg in  st.session_state['messages']:
    with st.chat_message(msg[0]):
        st.write(msg[1])
